import time

import torch

from models.arc_margin import ArcMarginModel
from models.embedder import FaceExpressionEmbedder
from utils import parse_args

if __name__ == '__main__':
    checkpoint = 'BEST_checkpoint.tar'
    print('loading {}...'.format(checkpoint))
    start = time.time()
    checkpoint = torch.load(checkpoint)
    print('elapsed {} sec'.format(time.time() - start))
    embedder = checkpoint['model'].module
    metric_fc = checkpoint['metric_fc'].module
    # print(model)
    # print(type(model))

    # model.eval()
    embedder_fn = 'embedder.pt'
    print('saving {}...'.format(embedder_fn))
    start = time.time()
    torch.save(embedder.state_dict(), embedder_fn)
    print('elapsed {} sec'.format(time.time() - start))

    print('loading {}...'.format(embedder_fn))
    start = time.time()
    model = FaceExpressionEmbedder()
    model.load_state_dict(torch.load(embedder_fn))
    print('elapsed {} sec'.format(time.time() - start))

    metric_fc_fn = 'metric_fc.pt'
    print('saving {}...'.format(metric_fc_fn))
    start = time.time()
    torch.save(metric_fc.state_dict(), metric_fc_fn)
    print('elapsed {} sec'.format(time.time() - start))

    print('loading {}...'.format(metric_fc_fn))
    start = time.time()
    args = parse_args()
    metric_fc = ArcMarginModel(args)
    metric_fc.load_state_dict(torch.load(metric_fc_fn))
    print('elapsed {} sec'.format(time.time() - start))
